# Copyright 2020 Google
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import os

import nbformat
from nbconvert import PythonExporter

TEMPLATE = """
{% extends 'python.tpl'%}

{%- block header -%}
# Copyright 2020 Google
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     https://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# ---                 This file has been autogenerated                    --- #
# ---              from docs/Readout-Data-Collection.ipynb                --- #
# ---                   Do not edit this file directly                    --- #
{% endblock header %}

## remove markdown cells
{%- block markdowncell -%}
{% endblock markdowncell %}

## remove execution count
{% block in_prompt %}
{% endblock in_prompt %}
"""

REPO_DIR = os.path.abspath(os.path.join(os.path.dirname(__file__), '..'))


def save_tasks_py(nb):
    """Select cells appropriate for `tasks.py`."""
    shouldnt_start_with = '# Put in a file named run-readout-scan.py'
    tasks_py = nb.copy()
    tasks_py.cells = [cell for cell in tasks_py.cells
                      if not cell.source.startswith(shouldnt_start_with)
                      and not cell.source.startswith('#@title')
                      and not any(l.strip().startswith('!') for l in cell.source.splitlines())]

    exporter = PythonExporter()
    exporter.raw_template = TEMPLATE
    output, resources = exporter.from_notebook_node(tasks_py)
    with open(f'{REPO_DIR}/recirq/readout_scan/tasks.py', 'w') as f:
        f.write(output)


def save_run_py(nb):
    """Select cells appropriate for `run-readout-scan.py`."""
    should_start_with = '# Put in a file named run-readout-scan.py'
    run_py = nb.copy()
    run_py.cells = [cell for cell in run_py.cells
                    if cell.source.startswith(should_start_with)]

    # Fix-up the content
    first_cell_lines = run_py.cells[0].source.splitlines()
    new_first_cell_lines = []
    added_imports = False
    for line in first_cell_lines:
        if line.startswith(should_start_with):
            # Ignore this comment, which is out of place in the exported py file
            continue

        if line.startswith('import') and not added_imports:
            # Add imports that aren't necessary when it's in one file.
            new_first_cell_lines += [
                'from recirq.readout_scan.tasks import ReadoutScanTask, run_readout_scan'
            ]
            added_imports = True

        # Rest of the lines are unmodified
        new_first_cell_lines += [line]

    # Mutate the original cell
    run_py.cells[0].source = '\n'.join(new_first_cell_lines)

    exporter = PythonExporter()
    exporter.raw_template = TEMPLATE
    output, resources = exporter.from_notebook_node(run_py)
    with open(f'{REPO_DIR}/recirq/readout_scan/run-readout-scan.py', 'w') as f:
        f.write(output)


def main():
    with open(f'{REPO_DIR}/docs/guide/data_collection.ipynb') as f:
        nb = nbformat.read(f, as_version=4)

    save_tasks_py(nb)
    save_run_py(nb)


if __name__ == '__main__':
    main()
